package bytecode

import (
	"context"
	"errors"

	"tvl.fyi/tvix/eval/value"
)

// RuntimeKey is the context key under which the runtime is expected
// to be available in a context.
type RuntimeKey struct{}

// runtime is an interface used by thunk evaluation to access the
// relevant runtime details and facilitate thunk evaluation. The
// runtime is expected to make itself available in the context using a
// [RuntimeKey].
type runtime interface {
	Evaluate(c *Chunk, captures []value.NixValue) (value.NixValue, error)
}

// Thunk is the public (opaque) representation of a thunk.
type Thunk struct {
	repr thunkRepr
}

func (t *Thunk) Force(ctx context.Context) (value.NixValue, error) {
	// TODO(tazjin): optimise already evaluated case?
	repr := t.repr
	t.repr = blackhole{}
	val, err := repr.force(ctx)
	if err != nil {
		return nil, err
	}

	t.repr = evaluated{val}
	return val, nil
}

func (t *Thunk) NixType() string {
	// TODO(tazjin): need a context for actual forcing?
	val, err := t.Force(context.TODO())
	if err != nil {
		// TODO(tazjin): need to be able to return error
		return "[" + err.Error() + "]"
	}

	return val.NixType()
}

func (t *Thunk) String() string {
	e, ok := t.repr.(evaluated)
	if ok {
		return e.v.String()
	}

	return "<CODE>"
}

func (t *Thunk) Strict(ctx context.Context) error {
	val, err := t.Force(ctx)
	if err != nil {
		return err
	}

	return val.Strict(ctx)
}

func (t *Thunk) Capture(captures []value.NixValue) error {
	suspended, ok := t.repr.(suspended)
	if !ok {
		return errors.New("tvix bug: Thunk::Capture() called on non-suspended thunk")
	}

	suspended.captures = captures
	t.repr = suspended // TODO(tazjin): necessary??
	return nil
}

// Value returns the value contained within an evaluated thunk. The value may be
// nil if the thunk was not evaluated yet.
func (t *Thunk) Value() value.NixValue {
	e, ok := t.repr.(evaluated)
	if ok {
		return e.v
	}

	return nil
}

// Constructor methods

func NewNativeThunk(f func() value.NixValue) *Thunk {
	return &Thunk{lazyValue(f)}
}

func NewEvaluatedThunk(v value.NixValue) *Thunk {
	return &Thunk{evaluated{v}}
}

func NewSuspendedThunk(c *Chunk) *Thunk {
	return &Thunk{suspended{c, nil}}
}

// Internal thunk representation details.

type thunkRepr interface {
	force(ctx context.Context) (value.NixValue, error)
}

// blackhole represents a thunk that is already under evaluation.
type blackhole struct{}

func (blackhole) force(ctx context.Context) (value.NixValue, error) {
	return nil, errors.New("infinite recursion detected")
}

// lazyValue is is a Nix value that is lazily computed by Go code
// (usually through sync.OnceValue).
type lazyValue func() value.NixValue

func (l lazyValue) force(ctx context.Context) (value.NixValue, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	return l(), nil
}

// evaluated holds a fully evaluated Nix value, which does not need to
// be forced anymore. This is used to memoize the thunk result.
type evaluated struct {
	v value.NixValue
}

func (e evaluated) force(context.Context) (value.NixValue, error) {
	return e.v, nil
}

// suspended holds an unevaluated chunk of Nix code. This is the most
// common thunk state.
type suspended struct {
	chunk    *Chunk
	captures []value.NixValue
}

func (s suspended) force(ctx context.Context) (value.NixValue, error) {
	r, ok := ctx.Value(RuntimeKey{}).(runtime)
	if !ok {
		return nil, errors.New("tvix bug: no runtime was supplied")
	}

	return r.Evaluate(s.chunk, s.captures)
}

type ThunkSet map[*Thunk]struct{}

// Seen checks whether the thunkset has seen this thunk. The thunkset has seen
// all thunks about which it has been asked.
func (s ThunkSet) Seen(t *Thunk) bool {
	_, seen := s[t]
	if !seen {
		s[t] = struct{}{}
	}

	return seen
}
